#include"Disney.h"

// https://seblagarde.wordpress.com/2013/04/29/memo-on-fresnel-equations/
//
// The Schlick Fresnel approximation is:
//
// R = R(0) + (1 - R(0)) (1 - cos theta)^5,
//
// where R(0) is the reflectance at normal indicence.
inline Float SchlickWeight(Float cosTheta) {
	Float m = clamp(1 - cosTheta, 0, 1);
	return (m * m) * (m * m) * m;
}

inline Float FrSchlick(Float R0, Float cosTheta) {
	return mix( R0, 1, SchlickWeight(cosTheta));
}

inline Spectrum FrSchlick(const Spectrum &R0, Float cosTheta) 
{
	return mix(R0, Spectrum(1.), SchlickWeight(cosTheta));
}

// For a dielectric, R(0) = (eta - 1)^2 / (eta + 1)^2, assuming we're coming from air.
inline Float SchlickR0FromEta(Float eta) { return sqr(eta - 1) / sqr(eta + 1); }

inline Float GTR1(Float cosTheta, Float alpha) {
	Float alpha2 = alpha * alpha;
	return (alpha2 - 1) /
		(Pi * std::log(alpha2) * (1 + (alpha2 - 1) * cosTheta * cosTheta));
}

// Smith masking shadowing term.
inline Float smithG_GGX(Float cosTheta, Float alpha) {
	Float alpha2 = alpha * alpha;
	Float cosTheta2 = cosTheta * cosTheta;
	return 1 / (cosTheta + sqrt(alpha2 + cosTheta2 - alpha2 * cosTheta2));
}


Spectrum DisneyDiffuse::f(P3 &light, P3 &view)
{
	Float FL = SchlickWeight(abs(light.z)),
		FV = SchlickWeight(abs(view.z));

	// Diffuse fresnel - go from 1 at normal incidence to .5 at grazing.
	// Burley 2015, eq (4).
	return color * InvPi * (1 - FL / 2) * (1 - FV / 2);
}

Spectrum DisneySubsurface::f(P3 &light, P3 &view)
{
	P3 wh = light + view;
	if (wh.x == 0 && wh.y == 0 && wh.z == 0) return Spectrum(0.);
	wh = wh.norm();
	Float cosThetaD = view.dot(wh);

	// Fss90 used to "flatten" retroreflection based on roughness
	Float Fss90 = cosThetaD * cosThetaD * roughness;
	Float FL = SchlickWeight(abs(light.z)),
		FV = SchlickWeight(abs(view.z));
	Float Fss = mix(1.0, Fss90,FL) * mix(1.0, Fss90,FV);
	// 1.25 scale is used to (roughly) preserve albedo
	Float ss =
		1.25f * (Fss * (1 / (abs(light.z) + abs(view.z)) - .5f) + .5f);

	return color * InvPi * ss;
}

Spectrum DisneyRetro::f(P3 &light, P3 &view)
{
	P3 wh = light + view;
	if (wh.x == 0 && wh.y == 0 && wh.z == 0) return Spectrum(0.);
	wh = wh.norm();
	Float cosThetaD = view.dot(wh);

	Float FL = SchlickWeight(abs(light.z)),
		FV = SchlickWeight(abs(view.z));
	Float Rr = 2 * roughness * cosThetaD * cosThetaD;

	// Burley 2015, eq (4).
	return color * InvPi * Rr * (FL + FV + FL * FV * (Rr - 1));
}

Spectrum DisneySheen::f(P3 &light, P3 &view)
{
	P3 wh = light + view;
	if (wh.x == 0 && wh.y == 0 && wh.z == 0) return Spectrum(0.);
	wh = wh.norm();
	Float cosThetaD = view.dot(wh);

	return color * SchlickWeight(cosThetaD);
}

Spectrum DisneyClearcoat::f(P3 &light, P3 &view)
{
	P3 wh = light + view;
	if (wh.x == 0 && wh.y == 0 && wh.z == 0) return Spectrum(0.);
	wh = wh.norm();
	// Clearcoat has ior = 1.5 hardcoded -> F0 = 0.04. It then uses the
	// GTR1 distribution, which has even fatter tails than Trowbridge-Reitz
	// (which is GTR2).
	Float Dr = GTR1(abs(wh.z), gloss);
	Float Fr = FrSchlick(.04, light.dot(wh));
	// The geometric term always based on alpha = 0.25.
	Float Gr =
		smithG_GGX(abs(light.z), .25) * smithG_GGX(abs(view.z), .25);

	return weight * Gr * Fr * Dr / 4;
}
P3 DisneyClearcoat::ImportanceSample(P3 &light, Float* rands, Float *pdf)
{
	if (light.z == 0) //no reflection
	{
		*pdf = 0;
		return P3();
	}

	Float alpha2 = gloss * gloss;
	Float cosTheta = sqrtf(
		max(Float(0), (1 - std::pow(alpha2, 1 - rands[0])) / (1 - alpha2)));
	Float sinTheta = sqrtf(max((Float)0, 1 - cosTheta * cosTheta));
	Float phi = 2 * Pi * rands[1];
	P3 wh(sinTheta*cos(phi), sinTheta*sin(phi), cosTheta);
	if (light.z*wh.z<=0) wh = -wh;

	P3 dir = -light.reflect(wh);
	if (light.z*dir.z <= 0)
	{
		*pdf = 0;
		return P3();
	}

	*pdf = getPdf(light,dir);
	return dir;
}
Float DisneyClearcoat::getPdf(P3 &light, P3 &view)
{
	if (light.z*view.z<=0) return 0;

	P3 wh = light + view;
	if (wh.x == 0 && wh.y == 0 && wh.z == 0) return 0;
	wh = wh.norm();

	// The sampling routine samples wh exactly from the GTR1 distribution.
	// Thus, the final value of the PDF is just the value of the
	// distribution for wh converted to a mesure with respect to the
	// surface normal.
	Float Dr = GTR1(abs(wh.z), gloss);
	return Dr * abs(wh.z) / (4 * light.dot(wh));
}

Spectrum DisneyMicroFacetReflection::f(P3 &light, P3 &view)
{
	Float cosThetaO = abs(light.z), cosThetaI = abs(view.z);
	P3 wh = light + view;
	// Handle degenerate cases for microfacet reflection
	if (cosThetaI == 0 || cosThetaO == 0) return Spectrum(0.f);
	if (wh.x == 0 && wh.y == 0 && wh.z == 0) return Spectrum(0.f);
	wh = wh.norm();
	// For the Fresnel call, make sure that wh is in the same hemisphere
	// as the surface normal, so that TIR is handled correctly.
	P3 wh_forward = wh.z < 0 ? -wh : wh;
	Spectrum F = getFresnel(view.dot(wh_forward));
	return facet_color * distribution->D(wh) * G(light, view) * F /
		(4 * cosThetaI * cosThetaO);
}
P3 DisneyMicroFacetReflection::ImportanceSample(P3 &light, Float* rands, Float *pdf)
{
	// Sample microfacet orientation $\wh$ and reflected direction $\wi$
	if (light.z == 0)
	{
		*pdf = 0;
		return P3();
	}
	P3 wh = distribution->Sample_wh(light, rands);
	if (light.dot(wh) < 0)   // Should be rare
	{
		*pdf = 0;
		return P3();
	}
	P3 dir = -light.reflect(wh);
	if (light.z*dir.z <= 0)
	{
		*pdf = 0;
		return P3();
	}

	// Compute PDF of _wi_ for microfacet reflection
	*pdf = distribution->distributionPdf(light, wh) / (4 * light.dot(wh));
	return dir;
}
Float DisneyMicroFacetReflection::getPdf(P3 &light, P3 &view)
{
	if (light.z*view.z <= 0) return 0;
	P3 wh = (light + view).norm();
	return distribution->distributionPdf(light, wh) / (4 * light.dot(wh));
}

Spectrum DisneyMicroFacetReflection::getFresnel(Float cosI)
{
	return mix(Spectrum(FrDielectric(cosI, 1, eta)),
		FrSchlick(fresnel_color, cosI), metallic);
}
inline Float DisneyMicroFacetReflection::RoughnessToAlpha(Float roughness)
{
	roughness = max(roughness, 0.001f);
	Float x = std::log(roughness), x2 = x * x, x4 = x2 * x2;
	return 1.62142f + 0.819955f * x + 0.1734f * x2 + 0.0171201f * x2*x +
		0.000640711f *x4;
}
inline Float DisneyMicroFacetReflection::G(P3 &wo, P3 &wi)
{
	// Disney uses the separable masking-shadowing model.
	return distribution->G1(wo) * distribution->G1(wi);
}

Spectrum DisneyMicroFacetTransmission::f(P3 &light, P3 &view)
{
	if (light.z*view.z > 0) return Spectrum(0);  // transmission only

	//P3 norm(0, 0, 1);
	//norm.z = light.z < 0 ? -1 : 1;
	//P3 dir;
	//if (light.z < 0)
	//	dir = (-light).refract(norm, etaB, etaA);
	//else
	//	dir = (-light).refract(norm, etaA, etaB);
	//bool into = light.z > 0;
	//Float a = etaB - 1, b = etaB + 1;
	//Float coef = a * a / (b*b), c = 1 - (into ? abs(light.z) : abs(dir.z));
	//Float Re = coef + (1 - coef)*c*c*c*c*c;
	//Float Tr = 1 - Re;
	//return Tr*color;

	Float cosThetaO = light.z;
	Float cosThetaI = view.z;
	if (cosThetaI == 0 || cosThetaO == 0) return Spectrum(0);

	Float eta = light.z > 0 ? (etaB / etaA) : (etaA / etaB);
	P3 wh = (light + view * eta).norm();
	if (wh.z < 0) wh = -wh;

	Float LH = light.dot(wh);
	Float VH = view.dot(wh);
	Spectrum F = getFresnelDielectric(LH);

	Float sqrtDenom = LH + eta * VH;
	Float factor = (mode == TransportMode::Radiance) ? (1 / eta) : 1;

	Spectrum ret_color= (Spectrum(1.f) - F) * color *
		std::abs(distribution->D(wh) * distribution->G(light, view) * eta * eta *
			abs(VH) * abs(LH) * factor * factor /
			(cosThetaI * cosThetaO * sqrtDenom * sqrtDenom));
	return ret_color;
}
P3 DisneyMicroFacetTransmission::ImportanceSample(P3 &light, Float* rands, Float *pdf)
{
	if (light.z == 0)
	{
		*pdf = 0;
		return P3();
	}
	P3 wh = distribution->Sample_wh(light, rands);
	if (light.dot(wh) < 0)
	{
		*pdf = 0;
		return P3();
	}
	//P3 norm(0, 0, 1);
	//norm.z = light.z < 0 ? -1 : 1;
	P3 dir;
	if (light.z < 0)
		dir = (light).refract(wh, etaB, etaA);
	else
		dir = (light).refract(wh, etaA, etaB);
	if (dir.x == 0 && dir.y == 0 && dir.z == 0)
	{
		//dir = -light.reflect(wh);

		//*pdf = distribution->distributionPdf(light, wh) / (4 * light.dot(wh));
		//return dir;
		*pdf = 0;
		return dir;
	}
	*pdf = getPdf(light, dir);
	if (dir.z*light.z > 0)dir = -dir;
	return dir;
}
Float DisneyMicroFacetTransmission::getPdf(P3 &light, P3 &view)
{
	if (light.z*view.z>0) return 0;
	//return 1;
	// Compute $\wh$ from $\wo$ and $\wi$ for microfacet transmission
	Float eta = light.z > 0 ? (etaB / etaA) : (etaA / etaB);
	P3 wh = (light + view * eta).norm();

	// Compute change of variables _dwh\_dwi_ for microfacet transmission
	Float LH = light.dot(wh), VH = view.dot(wh);
	Float sqrtDenom = LH + eta * VH;
	Float dwh_dwi =
		std::abs((eta * eta * VH) / (sqrtDenom * sqrtDenom));
	return distribution->distributionPdf(light, wh) * dwh_dwi;
}

